# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
rStar-Math 搜索树基础模块

这个模块定义了搜索算法的基础树结构和代码执行工具。BaseTree 是所有搜索算法的基类，
提供了树形搜索的基本功能和代码执行环境。

MCTS 算法核心概念：
- Q 值：节点的平均价值，用于评估节点质量
- 访问次数：节点被访问的次数，用于平衡探索和利用
- PUCT 值：结合 Q 值和置信上界的选择标准
- 回传：将叶子节点的价值向根节点传播

主要功能：
1. 定义搜索树的基础结构和接口
2. 提供 Python 代码执行环境
3. 实现代码提取和执行的工具函数
4. 管理搜索过程中的状态收集

核心组件：
- BaseTree: 搜索树基类，定义通用接口
- Python 解释器工具：安全执行推理过程中的代码
- 代码提取函数：从文本中提取可执行代码
- 超时保护：防止代码执行无限等待
"""
from __future__ import annotations  # 启用延迟注解，支持前向引用类型提示
import os  # 操作系统接口，用于文件路径验证和目录检查
from abc import abstractmethod  # 抽象基类装饰器，定义必须被子类实现的抽象方法
from termcolor import colored  # 终端彩色输出库，用于调试和日志的可视化显示
from typing import Optional, Any, Dict, List, Callable, Type, Tuple, Union  # 类型提示工具，提供静态类型检查支持
from pydantic import BaseModel, PrivateAttr, conlist, ConfigDict, field_validator  # 数据验证和序列化框架，确保配置和数据的类型安全
from omegaconf import DictConfig, OmegaConf  # 配置管理库，处理YAML配置文件和参数验证
from timeout_decorator import timeout  # 超时装饰器，为代码执行提供超时保护机制
from rstar_deepthink.config import BaseConfig  # 项目基础配置类，定义算法参数和模型配置
from rstar_deepthink.nodes.base_node import BaseNode  # 搜索树节点基类，表示MCTS和束搜索中的状态节点
from rstar_deepthink.tools.python_tool import PythonInterpreter  # Python代码执行器，安全执行数学推理过程中的计算代码
from rstar_deepthink.constants import TIMEOUT_SECONDS, TIMEOUT_MESSAGE, CODE_END, OUTPUT_END, CODE, ANSWER  # 项目常量定义，包含超时设置和文本解析标记


def _python_ast_init():
    """
    初始化 Python 解释器实例
    
    创建一个新的 Python 解释器，用于执行推理过程中生成的代码。
    解释器使用当前的全局命名空间，但局部命名空间为空。
    
    返回:
        PythonInterpreter: 配置好的 Python 解释器实例
    """
    python = PythonInterpreter(globals=globals(), locals=None)
    return python


def tool_wrapper(tool):
    """
    工具包装器：将工具对象包装为可调用函数
    
    将工具对象的 run 方法包装为简单的函数接口，
    便于在代码执行过程中统一调用。
    
    参数:
        tool: 具有 run 方法的工具对象
        
    返回:
        function: 包装后的工具函数
    """
    def _tool(query):
        return tool.run(query)
    return _tool


def no_action_wrapper(tool):
    """
    空操作包装器：用于处理无操作情况
    
    当没有具体操作需要执行时，返回提示信息让推理过程继续。
    这用于处理某些特殊情况下的占位操作。
    
    参数:
        tool: 工具对象（在此场景下不使用）
        
    返回:
        function: 返回提示信息的函数
    """
    def _tool(query):
        return "No action, no observation. Please continue to solve."
    return _tool


# 可用工具字典：映射工具名称到对应的执行函数
tools = {
    "python_interpreter": tool_wrapper(_python_ast_init()),  # Python 代码执行器
    "None": no_action_wrapper(_python_ast_init()),           # 空操作处理器
}


class BaseTree(BaseModel):
    """
    搜索树基类
    
    定义了所有搜索算法的基础接口和通用功能。作为 MCTS、束搜索等具体算法的基类，
    提供树形结构管理、节点创建、状态收集等核心功能。
    
    属性说明：
        config: 算法配置对象，包含模型路径、参数等
        question: 待解决的问题文本
        ground_truth: 问题的标准答案（可选）
        llm: 语言模型实例
        root: 搜索树的根节点
        current_node: 当前正在处理的节点
        stop: 生成停止标记列表
        node_max_retry: 节点最大重试次数
    
    核心职责：
    1. 管理搜索树的基本结构
    2. 提供节点创建的抽象接口
    3. 收集和组织搜索过程中的状态信息
    4. 验证配置参数的有效性
    """

    # 核心配置和问题信息
    config: Any                                        # 算法配置对象
    question: str                                      # 待解决的问题
    ground_truth: Optional[Union[str, List[str]]] = None  # 标准答案
    
    # 模型和搜索状态
    llm: Any = None                                    # 语言模型实例
    root: Optional[Type[BaseNode]] = None              # 根节点
    current_node: Optional[Type[BaseNode]] = None      # 当前节点
    
    # 控制参数
    stop: Optional[List[str]] = None                   # 停止标记
    node_max_retry: int = 5                           # 最大重试次数
    
    def __init__(self, **kwargs) -> None:
        """
        初始化搜索树
        
        设置基本配置，创建根节点，并初始化搜索状态。
        处理停止标记的配置转换（从 OmegaConf 格式到标准列表）。
        
        参数:
            **kwargs: 传递给父类的初始化参数
        """
        super().__init__(**kwargs)

        # 处理停止标记配置
        if self.config.stop:
            # 将 omegaconf.listconfig.ListConfig 转换为标准 list
            self.stop = OmegaConf.to_object(self.config.stop)

        # 创建根节点并设置为当前节点
        self.root = self.create_root()
        self.current_node = self.root
    
    @field_validator("config")
    def validate_config(cls, cfg: Any):
        """
        验证配置对象的有效性
        
        检查配置对象的类型和模型目录的存在性，确保算法能够正常运行。
        对于 OpenAI API 模式，跳过文件系统路径检查。
        
        参数:
            cfg: 待验证的配置对象
            
        返回:
            Any: 验证通过的配置对象
            
        异常:
            ValueError: 模型目录不存在
            TypeError: 配置对象类型错误
        """
        if issubclass(type(cfg), DictConfig):
            # 对于 OpenAI API 模式，跳过文件系统路径检查
            if cfg.backend_type == "openai":
                return cfg
            # 对于 vLLM 模式，检查模型目录是否存在
            if not os.path.exists(cfg.model_dir):
                raise ValueError(f"Model directory \"{cfg.model_dir}\" cannot be found.")
            return cfg

        raise TypeError("Wrong type for `config`, must be subclass of BaseConfig")
    
    def create_root(self) -> Type[BaseNode]:
        """
        创建搜索树的根节点
        
        调用抽象方法 create_node 创建根节点，并在其状态中记录问题信息。
        根节点是整个搜索过程的起点。
        
        返回:
            Type[BaseNode]: 初始化完成的根节点
        """
        root = self.create_node()
        root.state["extra_info"] = f"question: {self.question}"
        return root

    @abstractmethod
    def create_node(self, parent: Optional[Type[BaseNode]] = None) -> Type[BaseNode]:
        """
        创建新节点的抽象方法
        
        子类必须实现此方法，定义具体的节点创建逻辑。
        不同的搜索算法可能需要不同类型的节点（如 MCTSNode）。
        
        参数:
            parent: 父节点，根节点时为 None
            
        返回:
            Type[BaseNode]: 新创建的节点实例
        """
        pass
    
    def collect_partial_solution(self, node: Type[BaseNode]) -> str:
        """
        收集从根节点到指定节点的部分解决方案
        
        从指定节点开始向根节点回溯，收集路径上所有节点的文本内容，
        然后按正确顺序拼接成完整的推理轨迹。
        
        参数:
            node: 目标叶子节点
            
        返回:
            str: 拼接后的完整推理轨迹
        """
        # 从叶子节点向根节点收集文本
        trajectory = []
        while node:
            if node.state['text']:
                trajectory.append(node.state['text'])
            node = node.parent
        # 反转顺序，从根到叶的正确顺序
        return "".join(reversed(trajectory))
    
    def return_states(self) -> Dict[str, Dict[str, str]]:
        """
        返回搜索树中所有节点的状态信息
        
        使用广度优先搜索遍历整个搜索树，收集每个节点的状态字典。
        这对于调试、分析和可视化搜索过程非常有用。
        
        返回:
            Dict[str, Dict[str, str]]: 节点标签到状态字典的映射
        """
        candidates = [self.root]  # 待处理节点队列
        states = {}               # 状态收集字典
        
        # 广度优先遍历
        while candidates:
            node = candidates.pop(0)
            states[node.tag] = node.state
            # 将子节点加入队列
            if node.has_children():
                candidates.extend(node.children)
        return states
    
def extract_program(result: str, last_only=False):
    """
    从文本中提取 Python 代码程序
    
    解析包含代码块标记的文本，提取其中的可执行 Python 代码。
    支持提取最后一个代码块或所有代码块。
    
    参数:
        result: 包含代码块的原始文本
        last_only: 是否只提取最后一个代码块，默认 False
        
    返回:
        str: 提取出的 Python 代码，多个代码块用分隔符连接
        
    代码块格式:
        <code>
        # Python 代码
        print("Hello")
        <end_of_code>
    """
    program = ""
    start = False
    # 移除步骤结束标记
    result = result.replace("<end_of_step>", "")
    
    for line in result.split("\n"):
        if line.find("<code>") != -1:
            if last_only:
                program = ""  # 只保留最后一个程序
            else:
                program += "\n# ========\n"  # 代码块分隔符
            start = True
        elif line.find("<end_of_code>") != -1:
            start = False
        elif start:
            program += line + "\n"
    
    # 如果没有找到代码块标记，可能整个输出都是程序
    if not program:
        program = result
    return program.strip()

def code_execution(
    node: Type[BaseNode], 
    parser_result: Dict[str, str],
) -> str:
    """
    执行解析后的代码并返回观察结果
    
    根据解析结果执行相应的工具操作，主要用于执行 Python 代码。
    包含超时保护和异常处理，确保代码执行的安全性。
    
    参数:
        node: 当前节点，用于收集历史代码输入
        parser_result: 解析后的操作结果，包含 action 和 action_input
        
    返回:
        str: 代码执行的观察结果或错误信息
        
    工作流程:
    1. 根据 action 选择对应的工具函数
    2. 收集历史代码输入以保持上下文
    3. 提取并执行当前代码片段
    4. 返回执行结果或异常信息
    """

    @timeout(TIMEOUT_SECONDS, use_signals=True, exception_message=TIMEOUT_MESSAGE)
    def _code_execution(node: Type[BaseNode], parser_result: Dict[str, str]) -> str:
        """
        内部代码执行函数（带超时保护）
        
        参数:
            node: 当前节点
            parser_result: 解析结果
            
        返回:
            str: 执行结果
        """
        # 根据动作类型选择工具
        action = parser_result["action"]
        tool_func = tools[action]

        # 收集历史代码输入以保持执行上下文
        history_action_inputs = collect_action_inputs(node, action)

        # 执行当前代码片段
        action_input = parser_result["action_input"]
        # 合并历史输入和当前输入，然后提取程序
        action_input = extract_program(''.join(history_action_inputs) + action_input)
        observation = str(tool_func(action_input)).strip()
        del tool_func  # 清理工具函数引用
        return observation
    
    try:
        observation = _code_execution(node, parser_result)
    except Exception as e:
        # 捕获所有异常并格式化为错误信息
        observation = "{}: {}".format(type(e).__name__, str(e))
    
    return observation


def collect_action_inputs(
    node: Type[BaseNode], 
    action: str,
) -> List[str]:
    """
    收集从根节点到当前节点路径上指定动作的所有输入
    
    沿着节点的父子关系向上遍历，收集所有匹配指定动作类型的输入。
    这用于保持代码执行的上下文连续性，避免重复执行已完成的代码。
    
    参数:
        node: 当前节点
        action: 要收集的动作类型（如 "python_interpreter"）
        
    返回:
        List[str]: 按时间顺序排列的动作输入列表
        
    工作流程:
    1. 从当前节点开始向上遍历到根节点
    2. 遇到输出结束或代码结束标记时停止收集
    3. 收集所有匹配动作类型且未超时的输入
    4. 反转列表以获得正确的时间顺序
    """
    action_inputs = []
    while node: 
        # 遇到输出结束或代码结束标记时停止收集
        if OUTPUT_END in node.state['text'] or CODE_END in node.state['text']:
            break
        # 检查节点是否有匹配的动作
        if node.state["action"] == action:
            action_input = node.state["action_input"]
            # 确保输入存在且没有超时错误
            if action_input and "TimeoutError" not in node.state["text"].split(action_input)[-1]:
                action_inputs.append(action_input)
        node = node.parent
    # 反转列表以获得正确的时间顺序（从早到晚）
    return action_inputs[::-1]


def code_run(solution):
    """
    运行解决方案中的代码并返回执行结果
    
    检查解决方案是否包含完整的代码结构标记，如果包含则提取并执行代码。
    这是一个独立的代码执行函数，不依赖于搜索树的节点结构。
    
    参数:
        solution: 包含代码的解决方案文本
        
    返回:
        str: 代码执行结果或原始解决方案（如果不包含代码）
        
    工作流程:
    1. 检查解决方案是否包含必要的代码标记
    2. 如果不包含完整标记，直接返回原始解决方案
    3. 提取代码并使用 Python 解释器执行
    4. 返回执行结果或异常信息
    """
    # 检查是否包含完整的代码结构标记
    if CODE not in solution or CODE_END not in solution or OUTPUT_END not in solution or ANSWER not in solution:
        return solution
    
    @timeout(TIMEOUT_SECONDS, use_signals=True, exception_message=TIMEOUT_MESSAGE)
    def _code_execution(solution: str) -> str:
        """
        内部代码执行函数（带超时保护）
        
        参数:
            solution: 解决方案文本
            
        返回:
            str: 执行结果
        """
        tool_func = tools['python_interpreter']
        action_input = extract_program(solution)
        observation = str(tool_func(action_input)).strip()
        del tool_func  # 清理工具函数引用
        return observation
    
    try:
        observation = _code_execution(solution)
    except Exception as e:
        # 捕获所有异常并格式化为错误信息
        observation = "{}: {}".format(type(e).__name__, str(e))
    
    return observation